package agent.componet.message;

import agent.componet.message.thread.ServiceThread;
import dev.langchain4j.data.message.*;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.springframework.stereotype.Component;

import java.util.LinkedList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;


@Component
public class ChatMessageManager extends ServiceThread {

    private static ChatMessageManager INSTANCE = new ChatMessageManager();

    public static ChatMessageManager getInstance(){return INSTANCE;}

    private ChatMemory chatMemory;

    private ChatMemoryHook Hook ;

    private final LinkedBlockingDeque<ChatMessageRequest> linkedBlockingDeque = new LinkedBlockingDeque<>();

    private AtomicBoolean flag = new AtomicBoolean(false);

    public void registeHook(ChatMemoryHook chatMemoryHook){
        if (flag.get()) {
            return;
        }

        if (flag.compareAndSet(false,true)) {
             Hook = chatMemoryHook;
             if (Hook == null) Hook = new ChatMemoryHook() {
                 @Override
                 public ChatMemory getChatMeory() {
                     return new DefaultChatMemory();
                 }
             };
             chatMemory = Hook.getChatMeory();
             super.start();
        }
    }

    // cas  没有双端检查锁

    @Override
    public String getServiceName() {
        return this.getClass().getName();
    }
    @Override
    public void run() {


    while (!this.isStopped())
    {
        try {
            ChatMessageRequest take = linkedBlockingDeque.take();
            //  如果是这样的话其实不会有并发 但是 在取出来的时候会有并发
            chatMemory.setId(take.getId());
            chatMemory.add(take.getChatMessage());

        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }

    }


    }
    // ==================== 用户消息（UserMessage）相关 ====================

    /**
     * 获取指定ID下的所有用户消息列表
     */
    public List<UserMessage> getUserMessageList(Integer id) {
        chatMemory.setId(id);
        return chatMemory.messages().stream()
                .filter(msg -> msg instanceof UserMessage)
                .map(msg -> (UserMessage) msg)
                .collect(Collectors.toList());
    }

    /**
     * 获取指定ID下的最新一条用户消息
     */
    public UserMessage getLatestUserMessage(Integer id) {
        List<UserMessage> userMessages = getUserMessageList(id);
        if (userMessages.isEmpty()) {
            return null;
        }
        // 返回最后一条消息（假设列表顺序为消息时间顺序）
        return userMessages.get(userMessages.size() - 1);
    }

    /**
     * 获取指定ID下的所有用户消息拼接成的字符串
     */
    public String getUserMessageString(Integer id) {
        return getUserMessageList(id).stream()
                .map(UserMessage::singleText)
                .collect(Collectors.joining("\n"));
    }

    // ==================== 系统消息（SystemMessage）相关 ====================

    /**
     * 获取指定ID下的所有系统消息列表
     */
    public List<SystemMessage> getSystemMessageList(Integer id) {
        chatMemory.setId(id);
        return chatMemory.messages().stream()
                .filter(msg -> msg instanceof SystemMessage)
                .map(msg -> (SystemMessage) msg)
                .collect(Collectors.toList());
    }

    /**
     * 获取指定ID下的最新一条系统消息
     */
    public SystemMessage getLatestSystemMessage(Integer id) {
        List<SystemMessage> systemMessages = getSystemMessageList(id);
        if (systemMessages.isEmpty()) {
            return null;
        }
        // 返回最后一条消息（假设列表顺序为消息时间顺序）
        return systemMessages.get(systemMessages.size() - 1);
    }

    /**
     * 获取指定ID下的所有系统消息拼接成的字符串
     */
    public String getSystemMessageString(Integer id) {
        return getSystemMessageList(id).stream()
                .map(SystemMessage::text)  // 假设SystemMessage也有singleText()方法
                .collect(Collectors.joining("\n"));
    }

    // ==================== 混合消息相关 ====================

    /**
     * 获取指定ID下的所有消息（包含用户消息和系统消息）
     */
    public List<ChatMessage> getAllMessageList(Integer id) {
        chatMemory.setId(id);
        return chatMemory.messages();
    }

    /**
     * 获取指定ID下的所有消息拼接成的字符串（带消息类型标识）
     */
    public String getAllMessageString(Integer id) {
        return getAllMessageList(id).stream()
                .map(msg -> {
                    if (msg instanceof UserMessage) {
                        return "用户：" + ((UserMessage) msg).singleText();
                    } else if (msg instanceof SystemMessage) {
                        return "系统：" + ((SystemMessage) msg).text();
                    }
                    return "未知：" + msg.toString();
                })
                .collect(Collectors.joining("\n"));
    }


    // ==================== 新增：获取最近N条消息（含用户和系统消息） ====================

    /**
     * 获取指定ID下最近的N条消息（包含用户消息和系统消息）
     * @param id 会话ID
     * @param count 要获取的消息数量
     * @return 最近count条消息列表（按时间顺序，最新的在最后）
     */
    public List<ChatMessage> getRecentMessages(Integer id, int count) {
        List<ChatMessage> allMessages = getAllMessageList(id);
        if (allMessages.size() <= count) {
            return allMessages;
        }
        // 截取从（总条数 - 要获取的条数）开始到最后的子列表
        return allMessages.subList(allMessages.size() - count, allMessages.size());
    }

    /**
     * 获取指定ID下最近的N条消息拼接成的字符串（带消息类型标识）
     * @param id 会话ID
     * @param count 要获取的消息数量
     * @return 格式化后的消息字符串
     */
    public String getRecentMessagesString(Integer id, int count) {
        return getRecentMessages(id, count).stream()
                .map(this::formatMessageWithType)
                .collect(Collectors.joining("\n"));
    }

    // 辅助方法：格式化消息带类型标识
    private String formatMessageWithType(ChatMessage msg) {
        if (msg instanceof UserMessage) {
            return "用户：" + ((UserMessage) msg).singleText();
        } else if (msg instanceof SystemMessage) {
            return "系统：" + ((SystemMessage) msg).text();
        }
        return "未知：" + msg.toString();
    }



    // 依赖的消息处理方法（假设这些方法在同一个类中或通过注入获得）
    private List<UserMessage> getUserMessage(Integer id) {
        chatMemory.setId(id);
        List<ChatMessage> messages = chatMemory.messages();
        return messages.stream()
                .filter(x -> x instanceof UserMessage)
                .map(x -> (UserMessage) x)
                .collect(Collectors.toList());
    }

    private UserMessage getUserMessageConverToSingle(Integer id) {
        List<UserMessage> messages = getUserMessage(id);
        String combinedContent = messages.stream()
                .map(UserMessage::singleText)
                .collect(Collectors.joining("\n"));
        return new UserMessage(combinedContent);
    }

    public void SendRequest(Integer id , ChatMessage chatMessage){
        if(id == null || chatMessage == null)return;
        ChatMessageRequest chatMessageRequest = new ChatMessageRequest(id, chatMessage);
        linkedBlockingDeque.add(chatMessageRequest);
    }

    @FunctionalInterface
    interface ChatMemoryHook {
    ChatMemory getChatMeory() ;
    }


    @Data
    @Builder
    static
    class ChatMessageRequest{
       private Integer id;
     private    ChatMessage chatMessage;
   }

  interface ChatMemory extends dev.langchain4j.memory.ChatMemory {
        void setId(Integer id);
}



@NoArgsConstructor
@AllArgsConstructor
@Data
static
class DefaultChatMemory implements ChatMemory{
      private final static Integer DEFAULT_MAXWINDOWS = 10;
      private final ThreadLocal<Integer> id = new ThreadLocal<>();
      private  Integer maxMessages = DEFAULT_MAXWINDOWS;
      private  ChatMemoryStore store;

      public void setMaxwindows(Integer maxwindows){
                  maxMessages = maxwindows;
      }

      @Override
      public Object id() {
          return id.get();
      }

      @Override
      public void add(ChatMessage message) {
          List<ChatMessage> messages = messages();
          if (message instanceof SystemMessage) {
              Optional<SystemMessage> systemMessage = findSystemMessage(messages);
              if (systemMessage.isPresent()) {
                  if (systemMessage.get().equals(message)) {
                      return; // do not add the same system message
                  } else {
                      messages.remove(systemMessage.get()); // need to replace existing system message
                  }
              }
          }
          messages.add(message);
          ensureCapacity(messages, maxMessages);
          store.updateMessages(id, messages);
      }

      private  Optional<SystemMessage> findSystemMessage(List<ChatMessage> messages) {
          return messages.stream()
                  .filter(message -> message instanceof SystemMessage)
                  .map(message -> (SystemMessage) message)
                  .findAny();
      }

      @Override
      public List<ChatMessage> messages() {
          List<ChatMessage> messages = new LinkedList<>(store.getMessages(id));
          ensureCapacity(messages, maxMessages);
          return messages;
      }

      private void ensureCapacity(List<ChatMessage> messages, int maxMessages) {
          while (messages.size() > maxMessages) {

              int messageToEvictIndex = 0;
              if (messages.get(0) instanceof SystemMessage) {
                  messageToEvictIndex = 1;
              }

              ChatMessage evictedMessage = messages.remove(messageToEvictIndex);

              if (evictedMessage instanceof AiMessage && ((AiMessage) evictedMessage).hasToolExecutionRequests()) {
                  while (messages.size() > messageToEvictIndex
                          && messages.get(messageToEvictIndex) instanceof ToolExecutionResultMessage) {
                      // Some LLMs (e.g. OpenAI) prohibit ToolExecutionResultMessage(s) without corresponding AiMessage,
                      // so we have to automatically evict orphan ToolExecutionResultMessage(s) if AiMessage was evicted
                      ChatMessage orphanToolExecutionResultMessage = messages.remove(messageToEvictIndex);
                  }
              }
          }
      }

      @Override
      public void clear() {
          store.deleteMessages(id);
      }

    @Override
    public void setId(Integer id) {
       this.id.set(id); ;
    }
}
}
